/**

 */

#include "mem_allocator_c.h"

#include <errno.h>
#include <fcntl.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdlib.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>
#include <stdio.h>
#include <string.h>

enum MyLogLevel {
  MYLOG_LEVEL_TRACE = 0,
  MYLOG_LEVEL_DEBUG = 1,
  MYLOG_LEVEL_INFO = 2,
  MYLOG_LEVEL_WARN = 3,
  MYLOG_LEVEL_ERROR = 4,
  MYLOG_LEVEL_CRITICAL = 5,
  MYLOG_LEVEL_OFF = 6
};

#define MY_LOG(level, ...) \
do { \
    char log_buf[1024] = {0}; \
    int log_len = snprintf(log_buf, sizeof(log_buf), __VA_ARGS__); \
    if (log_len > 0) { \
      log_buf[sizeof(log_buf) - 1] = '\0'; \
      printf("%s\n", log_buf); \
    } \
} while(0)


#define MY_LOG_DEBUG(...) MY_LOG(MYLOG_LEVEL_DEBUG, __VA_ARGS__)
#define MY_LOG_INFO(...) MY_LOG(MYLOG_LEVEL_INFO, __VA_ARGS__)
#define MY_LOG_WARN(...) MY_LOG(MYLOG_LEVEL_WARN, __VA_ARGS__)
#define MY_LOG_ERROR(...) MY_LOG(MYLOG_LEVEL_ERROR, __VA_ARGS__)
#define MY_LOG_FATAL(...) MY_LOG(MYLOG_LEVEL_FATAL, __VA_ARGS__)

#define MY_LOG_VERIFY(_condition, ...)          \
  do {                                            \
    if (!(_condition)) {                          \
      MY_LOG(MYLOG_LEVEL_FATAL, __VA_ARGS__); \
      abort();                                    \
    }                                             \
  } while (0)

typedef struct MemAllocator MMapMemAllocator;
typedef struct MemAllocator DefaultMemAllocator;

static void mmap_params_free(MMapParams* mmap_params) {
  if (mmap_params == NULL) {
    return;
  }
  MemRange* range_list = mmap_params->range_list;
  while (range_list != NULL) {
    MemRange* range = range_list;
    range_list = range->next;
    free(range);
  }
  mmap_params->range_list = NULL;
  free(mmap_params);
  if (mmap_params->map_base != MAP_FAILED) {
    munmap(mmap_params->map_base, mmap_params->high_addr - mmap_params->base_addr + 1);
  }

  if (mmap_params->mem_fd >= 0) {
    close(mmap_params->mem_fd);
    mmap_params->mem_fd = MEM_FD_INIT;
  }
}

static MemRange* search_first_applicable_range(const MMapMemAllocator* allocator, uint32_t size_need) {
  if (allocator == NULL) {
    return NULL;
  }

  MemRange* range = allocator->mmap_params->range_list;
  while (range != NULL) {
    if (!range->is_dirty && range->size >= size_need) {
      return range;
    }
    range = range->next;
  }
  return NULL;
}

static int8_t* alloc_in_free_range(MMapMemAllocator* allocator, MemRange* range, uint32_t size_need) {
  if (allocator == NULL || range == NULL) {
    return NULL;
  }

  MY_LOG_VERIFY(!range->is_dirty && size_need <= range->size, "Try to alloc memory in an inapplicable range!");

  if (range->size - size_need < allocator->mmap_params->page_size) {
    range->is_dirty = true;
    memset(range->offset, 0x0, range->size);
    return range->offset;
  }

  MemRange* alloc_range = (MemRange*)malloc(sizeof(MemRange));
  alloc_range->offset = range->offset;
  alloc_range->end = range->offset + size_need - 1;
  alloc_range->size = size_need;
  alloc_range->is_dirty = false;
  alloc_range->next = range;
  alloc_range->pre = range->pre;

  if (alloc_range->pre != NULL) {
    alloc_range->pre->next = alloc_range;
  } else {
    allocator->mmap_params->range_list = alloc_range;
  }

  range->offset = alloc_range->end + 1;
  range->size -= size_need;
  range->pre = alloc_range;

  return alloc_range->offset;
}

static MemRange* find_target_range(const MMapMemAllocator* allocator, int8_t* buffer) {
  if (allocator == NULL || buffer == NULL) {
    return NULL;
  }

  MemRange* range = allocator->mmap_params->range_list;
  while (range != NULL) {
    if (range->offset == buffer) {
      return range;
    }
    range = range->next;
  }
  return NULL;
}

static void merge_contiguous_range(MMapMemAllocator* allocator, MemRange* first, MemRange* second) {
  if (allocator == NULL || first == NULL || second == NULL) {
    return;
  }

  // merge second range into the first range
  first->end = second->end;
  first->size += second->size;

  // update the links in the range list
  if (second->next) {
    second->next->pre = first;
  }
  first->next = second->next;

  free(second);
}

static int8_t* mmap_allocator_calloc(MMapMemAllocator* allocator, uint32_t nmemb, uint32_t size) {
  MY_LOG_DEBUG("mmap_allocator_calloc\n");
  uint32_t size_need = nmemb * size;
  // align by page
  size_need += (allocator->mmap_params->page_size - 1);
  size_need &= ~(allocator->mmap_params->page_size - 1);
  if (size_need < allocator->mmap_params->page_size) {
    size_need = allocator->mmap_params->page_size;
  }

  // traversal range list to find the first applicable range
  MemRange* range = search_first_applicable_range(allocator, size_need);
  if (range == NULL) {
    MY_LOG_INFO("no range is applicable");
    return NULL;
  }

  MY_LOG_DEBUG("[MemMapAllocator] will alloc %d bytes at %p", size_need, range->offset);
  int8_t* mem_addr = alloc_in_free_range(allocator, range, size_need);

  MY_LOG_DEBUG("[MemMapAllocator] current range list:");
  range = allocator->mmap_params->range_list;
  while (range != NULL) {
    MY_LOG_DEBUG("[MemMapAllocator] <start:%p, end:%p, size:%d, %s>", range->offset, range->end, range->size,
                   range->is_dirty ? "dirty" : "free");
    range = range->next;
  }

  return mem_addr;
}

static void mmap_allocator_free(MemAllocator* allocator, int8_t* buffer) {
  if (allocator == NULL) {
    return;
  }

  // set a range's is_dirty flag to true.
  MemRange* range = find_target_range(allocator, buffer);
  MY_LOG_VERIFY(range, "try to free an invalid point %p", buffer);
  range->is_dirty = false;

  // merge ranges if the range before and/or after this range is free.
  MemRange* next_range = range->next;
  if (next_range && !next_range->is_dirty) {
    merge_contiguous_range(allocator, range, next_range);
  }

  MemRange* pre_range = range->pre;
  if (pre_range && !pre_range->is_dirty) {
    merge_contiguous_range(allocator, pre_range, range);
  }
}

static int8_t* default_allocator_calloc(MemAllocator* allocator, uint32_t nmemb, uint32_t size) {
  if (allocator == NULL) {
    return NULL;
  }
  MY_LOG_DEBUG("default_allocator_calloc\n");
  return (int8_t*)calloc(nmemb, size);
}

static void default_allocator_free(MemAllocator* allocator, int8_t* buffer) {
  if (allocator == NULL) {
    return;
  }
  MY_LOG_DEBUG("default_allocator_free\n");
  free(buffer);
}

static void init_mmap_allocator(MMapMemAllocator* allocator, uint32_t base_addr, uint32_t high_addr) {
  if (allocator == NULL) {
    return;
  }
  allocator->mmap_params = (MMapParams*)malloc(sizeof(MMapParams));
  MMapParams* params = allocator->mmap_params;
  params->base_addr = base_addr;
  params->high_addr = high_addr;
  params->map_base = MAP_FAILED;
  params->mem_fd = MEM_FD_INIT;
  params->page_size = sysconf(_SC_PAGESIZE);
  params->range_list = NULL;
  if (params->base_addr >= params->high_addr) {
    return;
  }

  params->mem_fd = open("/dev/zero", O_RDWR | O_NDELAY);
  if (params->mem_fd >= 0) {
    // params->map_base = (int8_t*)mmap(NULL, params->high_addr - params->base_addr + 1,
    //                         PROT_READ | PROT_WRITE, MAP_SHARED,
    //                         params->mem_fd, params->base_addr);

    params->map_base = (int8_t*)mmap((void*)params->base_addr, params->high_addr - params->base_addr + 1,
                                     PROT_READ | PROT_WRITE, MAP_SHARED | MAP_FIXED, params->mem_fd, 0);
    if (params->map_base != MAP_FAILED) {
      // init range list
      params->range_list = (MemRange*)malloc(sizeof(MemRange));
      params->range_list->offset = params->map_base;
      params->range_list->end = (int8_t*)params->map_base + (params->high_addr - params->base_addr + 1);
      params->range_list->size = params->high_addr - params->base_addr + 1;
      params->range_list->is_dirty = false;
      params->range_list->next = NULL;
      params->range_list->pre = NULL;
    } else {
      MY_LOG_WARN("map failed! errno:%d", errno);
    }
  } else {
    MY_LOG_WARN("open /dev/mem failed! errno:%d", errno);
  }

  memset(allocator->allocator_name, 0, sizeof(allocator->allocator_name));
  const char* const allocator_name = "mmap_allocator";
  strncpy(allocator->allocator_name, allocator_name, strlen(allocator_name));
  allocator->calloc = mmap_allocator_calloc;
  allocator->free = mmap_allocator_free;
  allocator->next = NULL;

  MY_LOG_INFO("init_mmap_allocator done");
}

static void init_default_allocator(DefaultMemAllocator* allocator) {
  memset(allocator->allocator_name, 0, sizeof(allocator->allocator_name));
  const char* const allocator_name = "default_allocator";
  strncpy(allocator->allocator_name, allocator_name, strlen(allocator_name));
  allocator->calloc = default_allocator_calloc;
  allocator->free = default_allocator_free;
  allocator->next = NULL;
  MY_LOG_INFO("init_default_allocator done");
}

void init_mem_delegate(MemAllocDelegate* delegate, uint32_t base_addr, uint32_t high_addr) {
  if (base_addr == 0) {
    base_addr = MMAP_BASE_ADDR;
  }
  if (high_addr == 0) {
    high_addr = MMAP_HIGH_ADDR;
  }
  MY_LOG_INFO("init_mem_delegate, base_addr: %x, high_addr: %x", base_addr, high_addr);

  MMapMemAllocator* mmap_allocator = NULL;
  DefaultMemAllocator* default_allocator = NULL;
#ifdef MEM_ALLOC_MAP
  mmap_allocator = (MMapMemAllocator*)malloc(sizeof(MMapMemAllocator));
#endif  // MEM_ALLOC_MAP

#ifdef MEM_ALLOC_DEFAULT
  default_allocator = (DefaultMemAllocator*)malloc(sizeof(DefaultMemAllocator));
#endif  // MEM_ALLOC_DEFAULT

  if (mmap_allocator != NULL) {
    init_mmap_allocator(mmap_allocator, base_addr, high_addr);
  }

  if (default_allocator != NULL) {
    init_default_allocator(default_allocator);
  }
  if (mmap_allocator != NULL) {
    delegate->allocator = mmap_allocator;
    if (default_allocator != NULL) {
      delegate->allocator->next = default_allocator;
    }
  } else if (default_allocator != NULL) {
    delegate->allocator = default_allocator;
  } else {
    MY_LOG_FATAL("delegate->allocator is NULL");
  }
  MY_LOG_INFO("init_mem_delegate done");
}

int8_t* mem_delegate_calloc(MemAllocDelegate* delegate, uint32_t nmemb, uint32_t size) {
  MY_LOG_INFO("mem_delegate_calloc, nmemb: %d, size: %d", nmemb, size);

  if (delegate->allocator == NULL || delegate->allocator->calloc == NULL) {
    return NULL;
  }
  return delegate->allocator->calloc(delegate->allocator, nmemb, size);
}

void mem_delegate_free(MemAllocDelegate* delegate, int8_t* buffer) {
  if (delegate->allocator == NULL || delegate->allocator->free == NULL) {
    return;
  }
  delegate->allocator->free(delegate->allocator, buffer);
  MY_LOG_INFO("mem_delegate_free done");
}

void destroy_mem_delegate(MemAllocDelegate* delegate) {
  MemAllocator* current = delegate->allocator;
  while (current != NULL) {
    printf("free allocator: %s\n", current->allocator_name);
    MemAllocator* next = current->next;
    if (current->mmap_params != NULL) {
      mmap_params_free(current->mmap_params);
    }
    free(current);
    current = next;
  }
  MY_LOG_INFO("destroy_mem_delegate done");
}
